"""Loader that loads word documents."""
import os
from abc import ABC
from typing import List

from docx import Document as DP
from docx2python import docx2python
from langchain.docstore.document import Document
from langchain.document_loaders.base import BaseLoader


class DocxLoader(BaseLoader, ABC):
    """Loads a DOCX .

    Defaults to check for local file, but if the file is a web path, it will download it
    to a temporary file, and use that, then clean up the temporary file after completion
    """

    def __init__(self, file_path: str, SENTENCE_SIZE=64):
        """Initialize with file path."""
        self.file_path = file_path
        if "~" in self.file_path:
            self.file_path = os.path.expanduser(self.file_path)

        self.SENTENCE_SIZE = SENTENCE_SIZE

    def load(self) -> List[Document]:
        TextList = []
        try:
            tmp_list = []

            # python-docx不能很好处理numbered list
            # docx_file = DP(self.file_path)
            # for paragraph in docx_file.paragraphs:
            #     if not paragraph.text.strip():
            #         continue
            #     tmp_list.append(paragraph.text.strip())

            docx_file = docx2python(self.file_path)
            for paragraph in docx_file.body[0][0][0]:
                if not paragraph.strip():
                    continue
                tmp_list.append(paragraph.strip())

            merge_list = []
            for text in tmp_list:
                # 非数字开头且长度较短，合并
                if merge_list and not text.startswith(r'\d') and len(text) <= 8 and len(
                        merge_list[-1]) <= self.SENTENCE_SIZE:
                    merge_list[-1] = merge_list[-1] + text
                    continue
                merge_list.append(text)

            for text in merge_list:
                TextList.append(Document(
                    page_content=text,
                    metadata={"source": self.file_path},
                ))

            tables = docx_file.tables
            for table in tables:
                for i, row in enumerate(table.rows):  # 读每行
                    row_content = []
                    for cell in row.cells:  # 读一行中的所有单元格
                        c = cell.text
                        if c.strip() and c not in row_content:
                            row_content.append(c)
                    row_text = ' '.join(row_content)
                    row_text = re.sub(r"\n", ' ', row_text)
                    if row_text.strip():
                        row_text = '<table> ' + row_text
                        TextList.append(Document(
                            page_content=row_text,
                            metadata={"source": self.file_path},
                        ))
        except Exception as e:
            print(e)
            print("Error: {}".format(self.file_path))
        return TextList


if __name__ == '__main__':
    loader = DocxLoader(
        r'/home/sky/project/project/tianma/8-24/langchain-glm/base_files/docx_file/IT部门知识库/TM-C-WI19001  第三方人员安全管理规定1.3.docx')
    res = loader.load()
    print()
